import os
import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from numpy import inf
from torch.utils.tensorboard import SummaryWriter

from replay_buffer import ReplayBuffer
from velodyne_env import GazeboEnv


'''
评估函数：
它在给定的若干回合中运行,计算模型的平均奖励和碰撞次数
动作是通过网络生成的,而环境通过env.step()函数执行该动作
最终输出是该模型在若干测试回合中的平均奖励和碰撞次数
'''
def evaluate(network, epoch, eval_episodes=10):
    avg_reward = 0.0
    col = 0  # 碰撞次数统计
    for _ in range(eval_episodes):
        count = 0
        state = env.reset()  # 重置环境，获取初始状态
        done = False
        while not done and count < 501:  # 直到环境结束或达到最大步数
            action = network.get_action(np.array(state))  # 获取模型输出的动作
            a_in = [(action[0] + 1) / 2, action[1]]  # 将动作转换为特定范围的输入
            state, reward, done, _ = env.step(a_in)  # 执行动作，获取下一状态和奖励
            avg_reward += reward
            count += 1
            if reward < -90:  # 碰撞检测，通常负奖励表示撞到障碍物
                col += 1
    avg_reward /= eval_episodes  # 计算平均奖励
    avg_col = col / eval_episodes  # 计算平均碰撞次数
    print("..............................................")
    print(
        "Average Reward over %i Evaluation Episodes, Epoch %i: %f, %f"
        % (eval_episodes, epoch, avg_reward, avg_col)
    )
    print("..............................................")
    return avg_reward  # 返回平均奖励


'''
Actor 类:
定义了策略网络(Actor),输入是状态向量,输出是动作
使用两层全连接网络(800、600个神经元),最后输出使用 Tanh 激活函数将动作值映射到[-1, 1]区间
'''
# 定义 Actor 网络，输入状态，输出动作
class Actor(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(Actor, self).__init__()

        # 第一层全连接：输入为状态维度，输出800个神经元
        self.layer_1 = nn.Linear(state_dim, 800)
        # 第二层全连接：输入800个神经元，输出600个神经元
        self.layer_2 = nn.Linear(800, 600)
        # 第三层全连接：输入600个神经元，输出为动作维度
        self.layer_3 = nn.Linear(600, action_dim)
        # Tanh激活函数：将输出动作限制在[-1, 1]之间
        self.tanh = nn.Tanh()

    def forward(self, s):
        # 使用ReLU激活函数计算隐藏层输出
        s = F.relu(self.layer_1(s))  # 通过第一层全连接，激活为ReLU
        s = F.relu(self.layer_2(s))  # 通过第二层全连接，激活为ReLU
        a = self.tanh(self.layer_3(s))  # 通过第三层全连接，激活为Tanh限制动作范围
        return a  # 返回动作输出

'''
Critic 类:
定义了两个Q值网络(Critic),分别估计状态-动作对的价值(Q值)
每个网络都通过全连接层计算，使用两个不同的状态和动作编码器
'''
# 定义 Critic 网络，输入为状态和动作，输出为两个Q值估计
class Critic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(Critic, self).__init__()

        # Q1网络的第一层：输入状态，输出800个神经元
        self.layer_1 = nn.Linear(state_dim, 800)
        # Q1网络的第二层：状态子层
        self.layer_2_s = nn.Linear(800, 600)
        # Q1网络的第二层：动作子层
        self.layer_2_a = nn.Linear(action_dim, 600)
        # Q1网络的第三层：输出单一Q值
        self.layer_3 = nn.Linear(600, 1)

        # Q2网络的第一层：输入状态，输出800个神经元
        self.layer_4 = nn.Linear(state_dim, 800)
        # Q2网络的第二层：状态子层
        self.layer_5_s = nn.Linear(800, 600)
        # Q2网络的第二层：动作子层
        self.layer_5_a = nn.Linear(action_dim, 600)
        # Q2网络的第三层：输出单一Q值
        self.layer_6 = nn.Linear(600, 1)

    def forward(self, s, a):
        # Q1网络的前向传播
        s1 = F.relu(self.layer_1(s))  # 状态通过第一层
        self.layer_2_s(s1)  # 状态子层计算
        self.layer_2_a(a)  # 动作子层计算
        s11 = torch.mm(s1, self.layer_2_s.weight.data.t())  # 手动矩阵乘法
        s12 = torch.mm(a, self.layer_2_a.weight.data.t())  # 手动矩阵乘法
        s1 = F.relu(s11 + s12 + self.layer_2_a.bias.data)  # 合并状态和动作子层的输出
        q1 = self.layer_3(s1)  # 输出Q1值

        # Q2网络的前向传播
        s2 = F.relu(self.layer_4(s))  # 状态通过Q2的第一层
        self.layer_5_s(s2)  # 状态子层计算
        self.layer_5_a(a)  # 动作子层计算
        s21 = torch.mm(s2, self.layer_5_s.weight.data.t())  # 手动矩阵乘法
        s22 = torch.mm(a, self.layer_5_a.weight.data.t())  # 手动矩阵乘法
        s2 = F.relu(s21 + s22 + self.layer_5_a.bias.data)  # 合并状态和动作子层的输出
        q2 = self.layer_6(s2)  # 输出Q2值
        return q1, q2  # 返回两个Q值估计

'''
Actor 和 Critic:
    分别用于策略生成(Actor)和评估策略表现(Critic);该类同时维护了目标网络 actor_target 和 critic_target,用于稳定训练
'''
# TD3 算法类，负责训练和策略生成
class TD3(object):
    
    def __init__(self, state_dim, action_dim, max_action):
        # 初始化 Actor 网络和其目标网络
        self.actor = Actor(state_dim, action_dim).to(device)  # 当前Actor网络
        self.actor_target = Actor(state_dim, action_dim).to(device)  # 目标Actor网络
        self.actor_target.load_state_dict(self.actor.state_dict())  # 将目标网络参数初始化为Actor网络参数
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters())  # 使用Adam优化器训练Actor网络

        # 初始化 Critic 网络和其目标网络
        self.critic = Critic(state_dim, action_dim).to(device)  # 当前Critic网络
        self.critic_target = Critic(state_dim, action_dim).to(device)  # 目标Critic网络
        self.critic_target.load_state_dict(self.critic.state_dict())  # 将目标网络参数初始化为Critic网络参数
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters())  # 使用Adam优化器训练Critic网络

        self.max_action = max_action  # 动作的最大值
        self.writer = SummaryWriter()  # 用于记录数据到TensorBoard
        self.iter_count = 0  # 用于记录迭代次数

    # 生成动作,根据当前状态使用Actor网络生成动作
    def get_action(self, state):
        state = torch.Tensor(state.reshape(1, -1)).to(device)  # 将状态转为PyTorch张量
        return self.actor(state).cpu().data.numpy().flatten()  # 使用Actor生成动作并返回为NumPy数组


    # TD3算法的训练周期函数，用于更新Actor和Critic网络的参数
    def train(
        self,
        replay_buffer,  # 经验回放缓冲区，用于采样训练数据
        iterations,  # 当前训练周期中的更新次数
        batch_size=100,  # 每次采样的批量大小
        discount=1,  # 折扣因子γ，默认值为1（无折扣），通常为0.99
        tau=0.005,  # 软更新参数τ，用于更新目标网络
        policy_noise=0.2,  # 用于动作噪声的标准差，避免策略陷入局部最优
        noise_clip=0.5,  # 限制噪声的范围，防止噪声过大影响训练
        policy_freq=2,  # 每隔policy_freq次Critic更新后更新一次Actor网络
    ):
        av_Q = 0  # 初始化平均Q值
        max_Q = -inf  # 初始化最大Q值
        av_loss = 0  # 初始化平均损失

        '''经验回放缓冲区采样:每次从经验回放缓冲区中随机采样一个批次的数据，包括状态、动作、奖励、done标志和下一个状态'''
        for it in range(iterations):  # 遍历每次迭代
            # 从经验回放缓冲区中采样一个批次的数据
            (
                batch_states,  # 当前状态
                batch_actions,  # 动作
                batch_rewards,  # 奖励
                batch_dones,  # 结束标志
                batch_next_states,  # 下一个状态
            ) = replay_buffer.sample_batch(batch_size)  # 从经验池中采样一个batch

            # (batch_states,batch_actions,batch_rewards,batch_dones,batch_next_states) = replay_buffer.sample_batch(batch_size)

            state = torch.Tensor(batch_states).to(device)  # 将状态转换为PyTorch张量并移动到GPU
            next_state = torch.Tensor(batch_next_states).to(device)  # 下一个状态
            action = torch.Tensor(batch_actions).to(device)  # 动作
            reward = torch.Tensor(batch_rewards).to(device)  # 奖励
            done = torch.Tensor(batch_dones).to(device)  # 结束标志

            '''目标Q值计算'''
            '''使用目标Actor网络从下一个状态生成动作，并添加噪声来增强探索性'''
            # 使用目标Actor网络根据next_state估算下一个动作
            next_action = self.actor_target(next_state)

            # 向估算的下一个动作添加噪声，以增强探索性
            noise = torch.Tensor(batch_actions).data.normal_(0, policy_noise).to(device)  # 正态噪声
            noise = noise.clamp(-noise_clip, noise_clip)  # 限制噪声的范围
            next_action = (next_action + noise).clamp(-self.max_action, self.max_action)  # 将动作限制在[-max_action, max_action]之间

            '''通过目标Critic网络估算下一个状态-动作对的Q值'''
            # 使用目标Critic网络计算Q值（target Q）: Q(s', a')
            target_Q1, target_Q2 = self.critic_target(next_state, next_action)  # 计算两个Q值
            target_Q = torch.min(target_Q1, target_Q2)  # 取两个Q值的较小值，减少过度估计

            '''通过Bellman方程计算目标Q值'''
            # 计算目标Q值：r + γ * (1 - done) * Q(s', a')
            av_Q += torch.mean(target_Q)  # 累加平均Q值
            max_Q = max(max_Q, torch.max(target_Q))  # 更新最大Q值
            target_Q = reward + ((1 - done) * discount * target_Q).detach()  # 根据Bellman方程计算目标Q值

            '''Critic网络更新'''
            '''使用当前Critic网络计算当前状态和动作的Q值'''
            # 使用当前Critic网络计算当前状态和动作的Q值：Q(s, a)
            current_Q1, current_Q2 = self.critic(state, action)

            '''使用均方误差(MSE)损失函数计算Critic网络的损失,并通过反向传播更新Critic网络的参数'''
            # 计算当前Q值与目标Q值之间的均方误差损失
            loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q)  # 损失函数为两个Q值的均方差之和

            # 执行梯度下降来更新Critic网络
            self.critic_optimizer.zero_grad()  # 清除上一次的梯度
            loss.backward()  # 反向传播计算梯度
            self.critic_optimizer.step()  # 更新Critic网络参数

            '''Actor网络更新'''
            '''每隔一定次数更新一次Actor网络,目的是最大化Critic网络估算的Q值'''
            '''使用策略梯度方法更新Actor网络的参数'''
            # 每隔policy_freq次Critic更新后更新一次Actor网络
            if it % policy_freq == 0:
                # 使用Critic网络对Actor网络输出的动作进行估计，计算梯度并反向传播
                # 更新策略网络，最大化Q值
                actor_grad, _ = self.critic(state, self.actor(state))  # 通过当前状态估算Q值
                actor_grad = -actor_grad.mean()  # 取负值表示最大化Q值（梯度上升）
                self.actor_optimizer.zero_grad()  # 清除上一次的梯度
                actor_grad.backward()  # 反向传播更新Actor的梯度
                self.actor_optimizer.step()  # 更新Actor网络参数

                '''通过软更新将当前Actor和Critic网络的参数逐渐复制到其对应的目标网络上，从而保证训练的稳定性'''
                # 使用软更新（soft update）更新目标Actor网络参数
                for param, target_param in zip(
                    self.actor.parameters(), self.actor_target.parameters()
                ):
                    target_param.data.copy_(
                        tau * param.data + (1 - tau) * target_param.data
                    )

                # 使用软更新（soft update）更新目标Critic网络参数
                for param, target_param in zip(
                    self.critic.parameters(), self.critic_target.parameters()
                ):
                    target_param.data.copy_(
                        tau * param.data + (1 - tau) * target_param.data
                    )

            av_loss += loss  # 累加损失

        '''TensorBoard记录'''
        '''将训练过程中的平均损失、平均Q值、最大Q值写入TensorBoard,方便后续的可视化分析'''
        self.iter_count += 1  # 记录迭代次数

        # 将当前训练周期的平均损失、平均Q值、最大Q值写入TensorBoard以便可视化
        self.writer.add_scalar("loss", av_loss / iterations, self.iter_count)  # 平均损失
        self.writer.add_scalar("Av. Q", av_Q / iterations, self.iter_count)  # 平均Q值
        self.writer.add_scalar("Max. Q", max_Q, self.iter_count)  # 最大Q值


    def save(self, filename, directory):
        torch.save(self.actor.state_dict(), "%s/%s_actor.pth" % (directory, filename))
        torch.save(self.critic.state_dict(), "%s/%s_critic.pth" % (directory, filename))

    def load(self, filename, directory):
        self.actor.load_state_dict(
            torch.load("%s/%s_actor.pth" % (directory, filename))
        )
        self.critic.load_state_dict(
            torch.load("%s/%s_critic.pth" % (directory, filename))
        )


# 选择设备，优先使用GPU (cuda) 如果不可用则使用CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 设置随机种子，用于控制随机数生成，使实验可重复
seed = 0

# 设置评估频率，表示每5000步执行一次评估
eval_freq = 5e3  # 5000步

# 每个episode的最大步数
max_ep = 500

# 评估时的episode数
eval_ep = 10

# 总的最大时间步数，表示训练过程中允许执行的最大步数
max_timesteps = 5e6  # 500万步

# 初始探索噪声的值，范围 [expl_min ... 1]，用于在开始时让模型尝试多样的动作
expl_noise = 1

# 探索噪声的衰减步数，表示噪声在500000步内线性衰减
expl_decay_steps = 500000

# 探索噪声的最小值，噪声在衰减完成后保持在0.1
expl_min = 0.1

# Mini-batch的大小，表示每次更新参数时使用多少样本
batch_size = 40

# 折扣因子，用于计算未来奖励的折现值，值越接近1越重视长期回报
discount = 0.99999

# Soft target更新的速率，决定目标网络参数更新的权重（TD3中的一种稳定机制）
tau = 0.005

# 给策略网络添加的噪声，用于增加探索性动作
policy_noise = 0.2

# 噪声的最大值，用于防止噪声过大导致训练不稳定
noise_clip = 0.5

# 策略网络（Actor）更新的频率，即每执行2次Critic更新才更新一次Actor
policy_freq = 2

# 重放缓冲区的最大大小，用于存储经历的经验（状态、动作、奖励等）
buffer_size = 1e6  # 100万条记录

# 存储模型时的文件名称
file_name = "TD3_velodyne"

# 是否保存模型
save_model = True

# 是否加载预训练模型
load_model = False

# 是否在靠近障碍物时采取随机动作
random_near_obstacle = True


# 创建保存结果的文件夹（如果不存在）
if not os.path.exists("./results"):
    os.makedirs("./results")

# 如果设置了保存模型，并且模型文件夹不存在，则创建文件夹
if save_model and not os.path.exists("./pytorch_models"):
    os.makedirs("./pytorch_models")


# 设置环境的状态维度和机器人状态维度
environment_dim = 20  # 环境的状态维度
robot_dim = 4  # 机器人自身状态维度

# 创建仿真环境，使用Gazebo模拟多个机器人场景，环境维度为environment_dim
env = GazeboEnv("multi_robot_scenario.launch", environment_dim)
# 等待5秒钟以确保环境加载完成
time.sleep(5)

# 设置随机种子，使训练过程可复现
torch.manual_seed(seed)
np.random.seed(seed)

# 定义状态维度和动作维度
state_dim = environment_dim + robot_dim  # 状态包括环境和机器人状态
action_dim = 2  # 动作维度为2（例如前进和转向）

# 定义动作的最大值
max_action = 1


# 创建TD3算法的神经网络
network = TD3(state_dim, action_dim, max_action)

# 创建经验重放缓冲区，用于存储训练过程中积累的经验
replay_buffer = ReplayBuffer(buffer_size, seed)

# 如果设置为加载预训练模型，则尝试加载模型参数
if load_model:
    try:
        # 加载预训练模型的参数
        network.load(file_name, "./pytorch_models")
    except:
        # 如果无法加载，则打印警告信息，并使用随机初始化的模型参数开始训练
        print("Could not load the stored model parameters, initializing training with random parameters")


# 创建用于存储评估结果的列表
evaluations = []

# 初始化时间步计数器
timestep = 0

# 初始化自上次评估以来的时间步数
timesteps_since_eval = 0

# 初始化episode计数
episode_num = 0

# 标记episode是否完成，初始状态为完成
done = True

# 初始化epoch计数
epoch = 1

# 计数随机动作的次数
count_rand_actions = 0
# 存储随机动作
random_action = []

# 开始训练循环，直到达到最大时间步
while timestep < max_timesteps:

    # 如果当前episode已经结束
    if done:
        # 如果不是第一次执行，进行网络训练
        if timestep != 0:
            network.train(
                replay_buffer,
                episode_timesteps,
                batch_size,
                discount,
                tau,
                policy_noise,
                noise_clip,
                policy_freq,
            )

        # 如果达到评估频率，则进行评估
        if timesteps_since_eval >= eval_freq:
            print("Validating")
            # 重置自上次评估以来的时间步数
            timesteps_since_eval %= eval_freq
            # 进行评估并将结果保存
            evaluations.append(
                evaluate(network=network, epoch=epoch, eval_episodes=eval_ep)
            )
            # 保存当前模型
            network.save(file_name, directory="./pytorch_models")
            # 保存评估结果
            np.save("./results/%s" % (file_name), evaluations)
            # 增加epoch计数
            epoch += 1

        # 重置环境并初始化状态
        state = env.reset()
        done = False

        # 初始化当前episode的奖励和时间步数
        episode_reward = 0
        episode_timesteps = 0
        episode_num += 1


    # 添加一些探索噪声
    if expl_noise > expl_min:
        expl_noise = expl_noise - ((1 - expl_min) / expl_decay_steps)

    # 从网络中获取动作
    action = network.get_action(np.array(state))
    # 添加探索噪声，并确保动作在有效范围内
    action = (action + np.random.normal(0, expl_noise, size=action_dim)).clip(
        -max_action, max_action
    )


    # 如果机器人靠近障碍物，则强制其采取随机动作以增加探索
    if random_near_obstacle:
        # 85%的概率判断是否需要执行随机动作
        if (
            np.random.uniform(0, 1) > 0.85
            and min(state[4:-8]) < 0.6
            and count_rand_actions < 1
        ):
            # 随机生成一个在8到15之间的值，用于决定随机动作的持续时间
            count_rand_actions = np.random.randint(8, 15)
            # 随机生成一个动作
            random_action = np.random.uniform(-1, 1, 2)

        # 如果仍然在执行随机动作
        if count_rand_actions > 0:
            count_rand_actions -= 1
            action = random_action  # 使用随机动作
            action[0] = -1  # 设置线速度为-1（向后移动）


    # 更新动作到线速度范围 [0, 1] 和角速度范围 [-1, 1]
    a_in = [(action[0] + 1) / 2, action[1]]

    # 执行动作并获得下一个状态、奖励、是否完成的标志和目标
    next_state, reward, done, target = env.step(a_in)
    done_bool = 0 if episode_timesteps + 1 == max_ep else int(done)
    done = 1 if episode_timesteps + 1 == max_ep else int(done)
    episode_reward += reward


    # 将当前的状态、动作、奖励和下一个状态添加到重放缓冲区
    replay_buffer.add(state, action, reward, done_bool, next_state)

    # 更新状态
    state = next_state
    # 更新当前episode的时间步数
    episode_timesteps += 1
    # 更新总的时间步数
    timestep += 1
    # 更新自上次评估以来的时间步数
    timesteps_since_eval += 1


# 训练结束后进行最后一次评估并保存
evaluations.append(evaluate(network=network, epoch=epoch, eval_episodes=eval_ep))
if save_model:
    network.save("%s" % file_name, directory="./models")
    
# 保存评估结果
np.save("./results/%s" % file_name, evaluations)
